import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image

# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# 定义模型结构（必须与训练时相同）
class DigitClassifier(nn.Module):
    def __init__(self):
        super(DigitClassifier, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(64 * 5 * 5, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 64 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# 加载预训练模型
model = DigitClassifier().to(device)
model.load_state_dict(torch.load('digit_ocr_model_pytorch.pth'))
model.eval()


# 图像预处理函数
def preprocess_image(image_path):
    img = Image.open(image_path).convert('L')  # 转为灰度图
    img = img.resize((28, 28))
    img = np.array(img)
    img = 255 - img  # 反转颜色（如果需要）
    img = img / 255.0  # 归一化
    img = torch.tensor(img, dtype=torch.float32).view(1, 1, 28, 28)
    return img


# 识别数字
def recognize_digit(image_path):
    processed_img = preprocess_image(image_path).to(device)
    with torch.no_grad():
        output = model(processed_img)
        probabilities = torch.nn.functional.softmax(output, dim=1)
        confidence, predicted = torch.max(probabilities, 1)


    return predicted.item(), confidence.item() * 100


# 使用示例
if __name__ == "__main__":
    test_image_path = 'test_digit.png'  # 替换为你的测试图像路径
    digit, confidence = recognize_digit(test_image_path)
    print(f"识别结果: {digit}, 置信度: {confidence:.2f}%")